#!/usr/bin/env python
# -*- coding: utf-8 -*-
"""
@Author: 邵奈一
@Email: shaonaiyi@163.com
@Date: 2024/11/15
@微信：shaonaiyi888
@微信公众号: 邵奈一 
"""
import torch
import torch.nn as nn
from four_four_build_network import MYNET
import torchvision.transforms as transforms
from torch.utils.data import DataLoader, Dataset
from four_two_load_data import classes_txt
from four_two_load_data import MyDataset

# 代码4-3
# 加载图像数据
# 首先将训练集和测试集文件途径和文件名以txt保存在一个文件夹中，路径自行定义
root = 'data4'  # 文件的储存位置
classes_txt(root + '/train', root + '/train.txt')
classes_txt(root + '/test', root + '/test.txt')

# 由于数据集图片尺寸不一，因此要进行resize（重设大小）
# 将图片大小重设为 64 * 64
device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
transform = transforms.Compose([transforms.Resize((64, 64)),
                                transforms.Grayscale(),
                                transforms.ToTensor()])

# 提取训练集和测试集图片的路径生成txt文件
# num_class 选取100种汉字  提出图片和标签
train_set = MyDataset(root + '/train.txt',
                      num_class=100,
                      transforms=transform)
test_set = MyDataset(root + '/test.txt',
                     num_class=100,
                     transforms=transform)
# 放入迭代器中
train_loader = DataLoader(train_set, batch_size=50, shuffle=True)
test_loader = DataLoader(test_set, batch_size=5473, shuffle=True)
# 这里的5473 是因为测试集为5973张图片，当进行迭代时取第二批500个图片进行测试
for step, (x, y) in enumerate(test_loader):
    test_x, labels_test = x.to(device), y.to(device)

# 创建网络实例
model = MYNET().to(device)

# 代码4-5
# 设置优化器
# 优化器
optimizer = torch.optim.Adam(model.parameters(), lr=0.001)

# 代码4-6
# 设置损失函数
loss_func = nn.CrossEntropyLoss()

# 代码4-7
# 训练网络
# 训练模型
EPOCH = 6
for epoch in range(EPOCH):
    for step, (x, y) in enumerate(train_loader):
        picture, labels = x.to(device), y.to(device)
        output = model(picture)
        loss = loss_func(output, labels)
        optimizer.zero_grad()
        loss.backward()
        optimizer.step()

        # 性能评估
        if step % 50 == 0:
            test_output = model(test_x)
            pred_y = torch.max(test_output, 1)[1].data.squeeze()
            accuracy = ((pred_y == labels_test).sum().item() /
                        labels_test.size(0))
            # 输出迭代次数、训练误差、测试准确率
            print('迭代次数:', epoch,
                  '| 训练损失:%.4f' % loss.data,
                  '| 测试准确率:', accuracy)

print('完成训练')

# 代码4-9
# 保存网络
# 保存模型
torch.save(model.state_dict(), 'tmp/4-model.pkl')